package no.ntnu.fp.net.co;

import java.io.EOFException;
import java.io.IOException;
import java.net.ConnectException;
import java.net.InetAddress;
import java.net.SocketException;
import java.net.SocketTimeoutException;
import java.net.UnknownHostException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Random;

import no.ntnu.fp.net.cl.ClException;
import no.ntnu.fp.net.cl.ClSocket;
import no.ntnu.fp.net.cl.KtnDatagram;
import no.ntnu.fp.net.cl.KtnDatagram.Flag;

/**
 * Implementation of the Connection-interface. <br>
 * <br>
 * This class implements the behaviour in the methods specified in the interface
 * {@link Connection} over the unreliable, connectionless network realised in
 * {@link ClSocket}. The base class, {@link AbstractConnection} implements some
 * of the functionality, leaving message passing and error handling to this
 * implementation.
 * 
 * @author Sebj�rn Birkeland and Stein Jakob Nordb�
 * @see no.ntnu.fp.net.co.Connection
 * @see no.ntnu.fp.net.cl.ClSocket
 */
public class ConnectionImpl extends AbstractConnection {

	/** Keeps track of the used ports for each server port. */
	private static Map<Integer, Boolean> usedPorts = Collections.synchronizedMap(new HashMap<Integer, Boolean>());

	/**
	 * Initialise initial sequence number and setup state machine.
	 * 
	 * @param myPort
	 *            - the local port to associate with this connection
	 */
	public ConnectionImpl(int myPort) {
		super();
		if (usedPorts.containsKey(myPort)) {
			int newPort = generateRandomPort();
			this.myPort = newPort;
		} else {
			this.myPort = myPort;
			usedPorts.put(myPort, true);
		}
		this.myAddress = this.getIPv4Address();
	}

	private String getIPv4Address() {
		try {
			return InetAddress.getLocalHost().getHostAddress();
		}
		catch (UnknownHostException e) {
			return "127.0.0.1";
		}
	}

	/**
	 * Establish a connection to a remote location.
	 * @param remoteAddress
	 *            - the remote IP-address to connect to
	 * @param remotePort
	 *            - the remote portnumber to connect to
	 * @throws IOException
	 *             If there's an I/O error.
	 * @throws java.net.SocketTimeoutException
	 *             If timeout expires before connection is completed.
	 * @see Connection#connect(InetAddress, int)
	 */
	public void connect(InetAddress remoteAddress, int remotePort) throws IOException, SocketTimeoutException {
		if (state != State.CLOSED)
			throw new IllegalStateException("Must be in CLOSED state to connect.");

		// Set address and port to connect to
		this.remoteAddress = remoteAddress.getHostAddress();
		this.remotePort = remotePort;

		// Do the 3-way handshake
		int attempts = 5;
		do {
			try {
				// Send SYN
				simplySendPacket(constructInternalPacket(Flag.SYN));
				state = State.SYN_SENT;
				// Receive SYN_ACK
				KtnDatagram ack = receiveAck();
				// Make sure SYN_ACK is valid
				if (!checkPacket(ack) || !checkFlags(ack, Flag.SYN_ACK)) {
					continue;
				}
				// Set new port allocated by server
				this.remotePort = ack.getSrc_port();
				// Send ACK
				sendAck(ack, false);
				// Go to ESTABLISHED
				state = State.ESTABLISHED;
				return;
				
				// Gotta catch them all
			} catch (ClException e) {
			} catch (ConnectException e) {
			} catch (SocketException e) {
			} catch (IOException e) {
			}
		} while (attempts-- > 0);
		
		// Did not complete the 3-way handshake
		state = State.CLOSED;
		throw new SocketTimeoutException("Error: Did not manage to connect after 5 attempts.");
	}

	/**
	 * Listen for, and accept, incoming connections.
	 * 
	 * @return A new ConnectionImpl-object representing the new connection.
	 * @see Connection#accept()
	 */
	public Connection accept() throws IOException, SocketTimeoutException {
		if (state != State.CLOSED)
			throw new IllegalStateException("Must be in CLOSED state to accept connection.");

		// Set state to LISTEN
		state = State.LISTEN;

		// Wait for valid SYN
		KtnDatagram syn = null;
		do {
			try {
				syn = receivePacket(true);
			} catch (Exception e) {
				// Ignore
			}
		} while (!checkPacket(syn) || !checkFlags(syn, Flag.SYN));

		// Create new Connection to return
		int port = generateRandomPort();
		ConnectionImpl connection = new ConnectionImpl(port);
		connection.state = State.SYN_RCVD;
		connection.remoteAddress = syn.getSrc_addr();
		connection.remotePort = syn.getSrc_port();

		// Reply SYN with SYN_ACK from the new connection
		try {
			connection.sendAck(syn, true);
		} catch (Exception e) {
			// Ignore
		}

		// Wait for ACK
		KtnDatagram ack = connection.receiveAck();
		// Check if ACK is valid
		if (!checkPacket(ack) || !checkFlags(ack, Flag.ACK)) {
			// No valid ACK received
			state = State.CLOSED;
			connection.state = State.CLOSED;
			return null;
		}

		// Finalize the connection hand-over
		connection.lastValidPacketReceived = ack;
		connection.state = State.ESTABLISHED;
		this.state = State.CLOSED;
		return connection;
	}


	/**
	 * Send a message from the application.
	 * 
	 * @param msg
	 *            - the String to be sent.
	 * @throws ConnectException
	 *             If no connection exists.
	 * @throws IOException
	 *             If no ACK was received.
	 * @see AbstractConnection#sendDataPacketWithRetransmit(KtnDatagram)
	 * @see no.ntnu.fp.net.co.Connection#send(String)
	 */
	public void send(String msg) throws ConnectException, IOException {
		if (state != State.ESTABLISHED)
			throw new ConnectException("Must be in ESTABLISHED state to send.");

		// Send package.
		KtnDatagram packet = constructDataPacket(msg), ack = null;
		// Try 3 times
		int attempts = 3;
		do {
			try {
				
				ack = sendDataPacketWithRetransmit(packet);
			} catch (IOException e) {
				// Ignore
			}
		} while (!checkPacket(ack) && attempts-- > 0);

		// Make sure ACK is valid
		if (checkPacket(ack) && checkFlags(ack, Flag.ACK)) {
			if (packet.getSeq_nr() == ack.getAck()) {
				// Packet sent successfully
				lastDataPacketSent = packet;
				lastValidPacketReceived = ack;
			} else {
				// Send packet again
				nextSequenceNo--;
				send(msg);
			}
		} else {
			// Error while sending.
			throw new IOException("Could not send. Did not receive valid ACK.");
		}
	}

    /**
     * Wait for incoming data.
     * 
     * @return The received data's payload as a String.
     * @see Connection#receive()
     * @see AbstractConnection#receivePacket(boolean)
     * @see AbstractConnection#sendAck(KtnDatagram, boolean)
     */	
	public String receive() throws ConnectException, IOException {
		return receive(Long.MAX_VALUE);
	}
	
	/**
	 * Same as receive() but with infinite wait time
	 */
	public String receive(long timeout) throws ConnectException, IOException {
		if (state != State.ESTABLISHED)
			throw new ConnectException("Must be in ESTABLISHED state to receive.");
		
		// Receive packet
		KtnDatagram packet = null;
		try {
			// Receive packet
			packet = receivePacket(false, timeout);
		} catch (EOFException e) {
			throw new EOFException();
		} catch (IOException e) {
			e.printStackTrace();
		}
		
		// Make sure packet is valid
		if (isValid(packet)) {
			try {
				// Reply with ACK
				// Sleep first
				try { Thread.sleep(200); } catch (InterruptedException e) {	}
				sendAck(packet, false);
				lastValidPacketReceived = packet;
				// Return payload
				return packet.toString();
			} catch (IOException e) {
				// System.err.println("Could not send ACK on received packet.");
				e.printStackTrace();
			}
		} 

		// Packet contained errors, send ACK on last valid packet
		try {
			sendAck(lastValidPacketReceived, false);
		} catch (IOException e) {
		}
		return receive();
	}

	/**
	 * Close the connection.
	 * @see Connection#close()
	 */
	public void close() throws IOException {
		// Server receives EOF
		if (disconnectRequest != null) {
			// Got FIN from client, initiate passive close
			// System.err.println("RECEIVED FIN FROM " + remoteAddress + ":" + remotePort);
			// Set last valid packet received to FIN
			lastValidPacketReceived = disconnectRequest;
			
			// Sleep
			try { Thread.sleep(50); } catch (InterruptedException e) { }
			
			// System.err.println("SENDING FIN_ACK TO " + remoteAddress + ":" + remotePort);
			// Send FIN_ACK
			try {
				sendAck(disconnectRequest, false);
				nextSequenceNo = disconnectRequest.getSeq_nr() + 1;
			} catch (IOException e) { }
			// System.err.println("FINACK SENT TO " + remoteAddress + ":" + remotePort);

			// Go to CLOSE_WAIT
			state = State.CLOSE_WAIT;
			
			// Sleep
			try { Thread.sleep(750); } catch (InterruptedException e) { }
			
			// Try to send FIN to client
			try {
				KtnDatagram fin = constructInternalPacket(Flag.FIN);
				// System.err.println("SENDING FIN TO " + remoteAddress + ":" + remotePort);
				simplySendPacket(fin);
				// System.err.println("FIN SENT TO " + remoteAddress + ":" + remotePort);
			} catch (ClException e) {
			} catch (IOException e) {
				throw new IOException("Error, could not send FIN to " + remoteAddress + ":" + remotePort);
			}
			
			// Go to LAST_ACK
			state = State.LAST_ACK;
			
			// Try to receive FINACK from client
			// System.err.println("WAITING FOR FINACK FROM " + remoteAddress + ":" + remotePort);
			KtnDatagram finAck = null;
			try {
				finAck = receiveAck();
				if (checkPacket(finAck) && checkFlags(finAck, Flag.ACK)) {
					// System.err.println("FINACK RECEIVED FROM " + remoteAddress + ":" + remotePort);
					state = State.CLOSED;
					// System.err.println(remoteAddress + ":" + remotePort + " IS CLOSED");
				} else {
					throw new IOException();
				}
			} catch (IOException e) {
				state = State.CLOSED;
				// System.err.println("SERVER IS CLOSED (DID NOT RECEIVE FINAL FINACK FROM CLIENT)");
			}
		} 
		
		// Client initiates close()
		else {
			// Sleep
			try { Thread.sleep(100); } catch (InterruptedException e) {	}
			
			// Try to send FIN to server
			try {
				KtnDatagram fin = constructInternalPacket(Flag.FIN);
				// System.err.println("SENDING FIN TO " + remoteAddress + ":" + remotePort);
				simplySendPacket(fin);
				// System.err.println("FIN SENT TO " + remoteAddress + ":" + remotePort);
			} catch (ClException e) {
				// Ignore A2.
			} catch (IOException e) {
				e.printStackTrace();
				throw new IOException("Error, could not send FIN to " + remoteAddress + ":" + remotePort);
			}
			
			// Go to FIN_WAIT_1
			state = State.FIN_WAIT_1;
			
			// Wait for FIN_ACK from server
			// System.err.println("WAITING FOR FINACK FROM " + remoteAddress + ":" + remotePort);
			try {
				KtnDatagram finAck = receiveAck();
				if (checkPacket(finAck) && checkFlags(finAck, Flag.ACK)) {
					// System.err.println("RECEIVED FINACK FROM  " + remoteAddress + ":" + remotePort);
					state = State.FIN_WAIT_2;
				} else {
					throw new IOException("Error, FIN from  " + remoteAddress + ":" + remotePort + " was not valid");
				}
			} catch (IOException e) {
				throw new IOException("Error, did not receive FIN_ACK from " + remoteAddress + ":" + remotePort);
			}
			
			// Sleep first
			try { Thread.sleep(500); } catch (InterruptedException e) { }
			
			// Wait for FIN from server
			// System.err.println("WAITING FOR FIN FROM " + remoteAddress + ":" + remotePort);
			try {
				// Receive FIN
				KtnDatagram fin = receivePacket(true);
				// System.err.println("RECEIVED FIN FROM " + remoteAddress + ":" + remotePort);
				if (!checkPacket(fin) || !checkFlags(fin, Flag.FIN)) {
					throw new IOException();
				}
				
				// Sleep to let other side get ready
				try { Thread.sleep(500); } catch (InterruptedException e) { }
				
				// Send ACK
				// System.err.println("SENDING FIN_ACK TO " + remoteAddress + ":" + remotePort);
				sendAck(fin, false);
				// System.err.println("FIN_ACK SENT TO " + remoteAddress + ":" + remotePort);
			} catch (IOException e) {
				// System.err.println("Error, did not get FIN or send FIN_ACK to server " + remoteAddress + ":" + remotePort);
				state = State.CLOSED;
				return;
			}
			// Go to TIME_WAIT
			state = State.TIME_WAIT;
			// System.err.println("CLIENT IS TIME_WAIT");
			
			// Go to CLOSED
			state = State.CLOSED;
			// System.err.println("CLIENT IS CLOSED");
			return;
		}
	}

	/**
	 * Test a packet for transmission errors. This function should only called
	 * with data or ACK packets in the ESTABLISHED state.
	 * 
	 * @param packet
	 *            Packet to test.
	 * @return true if packet is free of errors, false otherwise.
	 */
	protected boolean isValid(KtnDatagram packet) {
		boolean check = checkPacket(packet);
		boolean checksum = packet.calculateChecksum() == packet.getChecksum();
		boolean addressAndPort = checkAddress(packet) && checkPort(packet);
		boolean seqNr = packet.getSeq_nr() == lastValidPacketReceived.getSeq_nr() + 1;
		return check && checksum && addressAndPort && seqNr;
	}

	/**
	 * Generates a random new port and places it in the used ports list.
	 * @return A random new port number.
	 */
	private int generateRandomPort() {
		Random rand = new Random();
		int newPort = 0;
		while (true) {
			newPort = rand.nextInt(10000 - 1024 + 1) + 1024;
			if (!usedPorts.containsKey(newPort)) {
				usedPorts.put(newPort, true);
				break;
			}
		}
		return newPort;
	}

	/**
	 * Helper method for isValid(). Checks if the source address of the packet is the same as the remote address.
	 * @param packet
	 * @return true if source address in packet equals the remote address.
	 */
	private boolean checkAddress(KtnDatagram packet) {
		return packet.getSrc_addr().equals(remoteAddress);
	}

	/**
	 * Helper method for isValid(). Checks if the source port of the packet is the same as the remote port.
	 * @param packet
	 * @return true if source port in packet equals the remote port.
	 */
	private boolean checkPort(KtnDatagram packet) {
		return packet.getSrc_port() == remotePort;
	}
	
	/** Helper method to check flags */
	private boolean checkFlags(KtnDatagram packet, Flag ... flags) {
		for (Flag flag : flags) {
			if (packet.getFlag() != flag)
				return false;
		}
		return true;
	}
	
	/** Helper method to make sure packets are not null and have a correct checksum */
	private boolean checkPacket(KtnDatagram packet) {
		return packet != null;
	}

	// Getters for debugging.
	public String getLocalAddress() {
		return myAddress;
	}

	public String getRemoteAddress() {
		return remoteAddress;
	}

	public int getLocalPort() {
		return myPort;
	}

	public int getRemotePort() {
		return remotePort;
	}
}